import os
import random

import PIL.Image
import pandas as pd
from torch.utils.data import DataLoader
from torchvision import transforms

from flower_dataset import FlowerDataset


def get_target_trans():
    trans = {
        'train': transforms.Compose([
            lambda x: x - 1,
            transforms.ToTensor()
        ]),
        'val': lambda x: x - 1,
        'test': transforms.Compose([
            lambda x: x - 1,
            transforms.ToTensor()
        ])
    }

    return trans


def get_transform():
    trans = {
        'train': transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'val': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ]),
        'test': transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
    }

    return trans


def get_data(root_path: str, batch_size: int, transform=None, target_transform=None):
    # train_set = CIFAR10(
    #     root=root_path,
    #     train=True,
    #     transform=transform['train'] if transform is not None else None,
    #     target_transform=target_transform['train'] if target_transform is not None else None,
    #     download=True
    # )
    # # val_set = CIFAR10(
    # #     root=root_path,
    # #     split='val',
    # #     transform=transform['val'] if transform is not None else None,
    # #     target_transform=target_transform['val'] if target_transform is not None else None,
    # #     download=True
    # # )
    # test_set = CIFAR10(
    #     root=root_path,
    #     train=False,
    #     transform=transform['test'] if transform is not None else None,
    #     target_transform=target_transform['test'] if target_transform is not None else None,
    #     download=True
    # )
    train_set = FlowerDataset(
        root_path=root_path,
        type='train',
        transform=transform['train'] if transform is not None else None,
        target_transform=target_transform['train'] if target_transform is not None else None,
    )
    val_set = FlowerDataset(
        root_path=root_path,
        type='val',
        transform=transform['val'] if transform is not None else None,
        target_transform=target_transform['val'] if target_transform is not None else None,
    )
    test_set = FlowerDataset(
        root_path=root_path,
        type='test',
        transform=transform['test'] if transform is not None else None,
        target_transform=target_transform['test'] if target_transform is not None else None,
    )

    train_loader = DataLoader(
        train_set,
        batch_size=batch_size,
        shuffle=True
    )
    val_loader = DataLoader(
        val_set,
        batch_size=batch_size,
        shuffle=False
    )
    test_loader = DataLoader(
        test_set,
        batch_size=batch_size,
        shuffle=True
    )

    return train_set, val_set, test_set, train_loader, val_loader, test_loader
    # return train_set, test_set, train_loader, test_loader


def getSingleImage(root_path) -> (PIL.Image.Image, str):
    df = pd.read_csv(os.path.join(root_path, 'test.csv'))
    index = random.randint(0, len(df) - 1)
    data = df.loc[index]
    path, label = data['path'], data['label']
    img = PIL.Image.open(path)
    return img, label
